from Crypto.Cipher import PKCS1_OAEP
from Crypto.PublicKey import RSA
import base64


class RsaEnc:
    rsa_pubkey = None
    rsa_private = None
    pubkey = "************************"
    prikey = "*************"

    def __init__(self):
        try:
            self.rsa_pubkey = RSA.importKey(self.pubkey)
        except Exception as err:
            print('字符串密钥转rsa格式密钥错误', err)

        try:
            self.rsa_private = RSA.import_key(self.prikey)
        except Exception as err:
            print('字符串密钥转rsa格式密钥错误', err)

    def rsa_enc(self, data):
        try:
            cipher = PKCS1_OAEP.new(self.rsa_pubkey)
            return str(base64.b64encode(cipher.encrypt(data.encode("utf-8"))), "utf-8")
        except Exception as err:
            print('RSA加密失败', '', err)

    def rsa_dec(self, data):
        ret_data = b''
        try:
            cipher = PKCS1_OAEP.new(self.rsa_private)
            ret_data = cipher.decrypt(base64.b64decode(data))
        except Exception as err:
            print('RSA解密失败', '', err)
        return ret_data.decode("utf-8")

#
# def test():
#     r = RsaEnc()
#     print(r.rsa_pubkey)
#     print(r.rsa_private)
#     rst = r.rsa_enc("123456")  # 加密函数
#     print(rst)
#     print(r.rsa_dec(rst))  # 解密函数
